# -*- coding: utf-8 -*-
# @Time    : 2023/5/4 上午11:58
# @Author  : lilong
# @desc    :

import numpy as np

import tensorflow as tf
from keras.layers import *
import keras.backend as K


def seq_gather(x):
    seq, idxs = x

    batch_idxs = K.arange(0, K.shape(seq)[0])
    batch_idxs = K.expand_dims(batch_idxs, 1)

    idxs = K.cast(idxs, 'int32')
    idxs = K.concatenate([batch_idxs, idxs], 1)
    return tf.gather_nd(seq, idxs)


# 3x4
tokens_feature = np.array(
    # [[1, 2, 3, 4],
    #  [3, 4, 4, 5],
    #  [2, 3, 4, 5]]
    [[[1, 2, 3, 4],
     [3, 4, 4, 5],
     [2, 3, 4, 5]],
    [[1, 2, 3, 4],
     [3, 4, 4, 5],
     [2, 3, 4, 5]]]
)

# 1x3
head = np.array(
    [[1], [2]])

tail = np.array(
    [[0], [2]])

# sub_head_in = Input(shape=(1,))
# # sub_head_in()
# print(sub_head_in)
# tokens_feature = K.expand_dims(tokens_feature, 0)
# head = K.expand_dims(head, 0)
sub_head_feature = Lambda(seq_gather)([tf.convert_to_tensor(tokens_feature),
                                       tf.convert_to_tensor(head)])
sub_tail_feature = Lambda(seq_gather)([tf.convert_to_tensor(tokens_feature),
                                       tf.convert_to_tensor(tail)])

print(sub_head_feature)
sub_feature = Average()([sub_head_feature, sub_tail_feature])

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(sub_head_feature))
    print(sess.run(sub_tail_feature))
    print(sess.run(sub_feature))